import os
import shutil
from pathlib import Path

import pandas as pd
import random
from Final_assignment.tools import *

# 全局参数
train_percent = 0.8  # 训练集占比
train_label_csv = pd.DataFrame(columns=['filename', 'index', 'label'])  # 标签csv
test_label_csv = pd.DataFrame(columns=['filename', 'index', 'label'])  # 标签csv
train_index = 0  # 训练集索引
test_index = 0  # 测试集索引


# 划分数据集
def partition_dataset(cycle, train_folder, test_folder, stock_list):  # 周期，训练集文件夹，训练集文件夹,股票列表
    print("now is partition dataset ...")
    empty_folder(train_folder)  # 清空训练集文件夹
    empty_folder(test_folder)  # 清空测试集文件夹
    # 变量每一只股票
    for stock in stock_list:
        stock = stock_code_standardization(stock)   # 股票代码标准化
        copy_img_and_label(stock, cycle, train_folder, test_folder)
    print("训练集的数量：{},测试集的数量：{}".format(train_index, test_index))


# 清空文件夹中的内容
def empty_folder(folder_path):
    if os.path.exists(folder_path) is True:  # 若路径存在
        shutil.rmtree(folder_path)  # 删除文件夹和其中的内容
    os.mkdir(folder_path)  # 创建文件夹


# 准备训练集
def copy_img_and_label(stock_name, cycle, train_img_path_to, test_img_path_to):
    global train_index
    global test_index
    img_path_from = os.path.join("Kline_diagram", cycle + "_Kline", stock_name, )  # k线图文件夹
    label_path = os.path.join("Kline_diagram", cycle + "_label", stock_name + ".csv")  # 标记文件

    label_csv = pd.read_csv(label_path, index_col=0)  # 读取一只股票的标签数据

    img_nums = len(label_csv)  # 图片数量
    train_nums_divide = int(img_nums * train_percent)  # 训练集数量
    test_nums_divide = img_nums - train_nums_divide  # 测试集数量
    print("{}: 训练集应分得的图片数量：{},测试集应分得的图片数量：{}".format(stock_name, train_nums_divide, test_nums_divide))
    train_nums, test_nums = 0, 0
    index = 0
    while train_nums < train_nums_divide:  # 前面的给训练集
        copy_file(os.path.join(img_path_from,label_csv.iloc[index]["filename"])  , train_img_path_to)  # 分给训练集
        train_nums += 1
        train_index += 1
        train_label_csv.loc[train_index] = label_csv.iloc[index]
        index += 1
    while test_nums < test_nums_divide:  # 后面的给测试集
        copy_file(os.path.join(img_path_from,label_csv.iloc[index]["filename"]), test_img_path_to)  # 分给训练集
        test_nums += 1
        test_index += 1
        test_label_csv.loc[test_index] = label_csv.iloc[index]
        index += 1
    train_label_csv.to_csv(os.path.join(train_img_path_to, train_img_path_to + '.csv'))  # 保存训练集标签数据
    test_label_csv.to_csv(os.path.join(test_img_path_to, test_img_path_to + '.csv'))  # 保存测试集标签数据


# 复制函数,将文件srcfile复制到文件夹dstpath中
def copy_file(src_file, dst_path):
    if not os.path.isfile(src_file):
        print("%s not exist!" % (src_file))
    else:
        fpath, fname = os.path.split(src_file)  # 分离文件名和路径
        if not os.path.exists(dst_path):
            os.makedirs(dst_path)  # 创建路径
        shutil.copy(src_file, os.path.join(dst_path, fname))  # 复制文件


'''
pandas 文件读取csv时Unamed列解决办法
https://blog.csdn.net/weixin_43593330/article/details/91863692 
'''
